import os

base_dir = 'G:\\python\\keshe\\train2'
# # 训练、验证、测试数据集的目录
train_dir = os.path.join(base_dir, 'train')
validation_dir = os.path.join(base_dir, 'validation')
test_dir = os.path.join(base_dir, 'test')
# 猫训练图片所在目录
train_cats_dir = os.path.join(train_dir, 'cats')
# 狗训练图片所在目录
train_dogs_dir = os.path.join(train_dir, 'dogs')
# 猫验证图片所在目录
validation_cats_dir = os.path.join(validation_dir, 'cats')
# 狗验证数据集所在目录
validation_dogs_dir = os.path.join(validation_dir, 'dogs')
# 猫测试数据集所在目录
test_cats_dir = os.path.join(test_dir, 'cats')
# 狗测试数据集所在目录
test_dogs_dir = os.path.join(test_dir, 'dogs')


import keras
from keras import layers, models
from keras import optimizers, losses

model = models.Sequential()

model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(64, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Conv2D(128, (3, 3), activation='relu'))
model.add(layers.MaxPooling2D((2, 2)))
model.add(layers.Flatten())
model.add(layers.Dense(512, activation='relu'))
model.add(layers.Dense(1, activation='sigmoid'))
# 在编译之前可以查看网络的架构
model.summary()

# 编译模型
# model.compile(loss=losses.binary_crossentropy, optimizer=optimizers.RMSprop(lr=1e-4), metrics=['acc'])  # 这个是书上的，指定了学习率
model.compile(loss=losses.binary_crossentropy, optimizer='rmsprop', metrics=['acc'])

# keras.preprocessing.image图像处理辅助工具模块
from keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(rescale=1./255)
test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
                    train_dir,
                    target_size=(150, 150),
                    batch_size=20,
                    class_mode='binary')
validation_generator = train_datagen.flow_from_directory(
                    validation_dir,
                    target_size=(150, 150),
                    batch_size=20,
                    class_mode='binary')

# 可以看一下其中一个生成器的输出，生成的是150 * 150 的RGB图像(形状为（20， 150， 150， 3）)与二进制标签（形状为（20， ））组成的批量。
for data_batch, labels_batch in train_generator:
    print('data batch shape:', data_batch.shape)
    print('labels batch shape:', labels_batch.shape)
    break

# 利用批量生成器拟合模型
generator = model.fit_generator(
    train_generator,
    steps_per_epoch=100,
    epochs=10,
    validation_data=validation_generator,
    validation_steps=50)



model.save('cats_and_dogs_small_1.h5')  # 保存模型

# 绘制损失曲线和精度曲线
# 绘制损失曲线和精度曲线
import matplotlib.pyplot as plt

acc = generator.history.get('acc')
val_acc = generator.history.get('val_acc')
loss = generator.history.get('loss')
val_loss = generator.history.get('val_loss')

epochs = range(1, len(acc) + 1)

plt.xlabel('epochs')
plt.ylabel('percentage')
plt.plot(epochs, acc, 'r', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('Training and validation accuracy')
plt.legend()

plt.figure()

plt.xlabel('epochs')
plt.ylabel('percentage')
plt.plot(epochs, loss, 'r', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss')
plt.legend()

plt.show()
